from ultralytics import YOLO

if __name__ == '__main__':
    # 加载模型 - 使用预训练权重
    model = YOLO("runs/detect/train21/weights/best.pt")

    # 训练配置
    results = model.train(
        data="ultralytics/cfg/datasets/helmet3.yaml",  # 确保使用包含增强参数的配置文件
        epochs=50,
        freeze=10,  # 减少冻结层数，加速后期收敛
        lr0=0.005,  # 恢复标准学习率
        lrf=0.001,  # 最终学习率比例
        warmup_epochs=3,  # 添加学习率预热
        warmup_momentum=0.8,
        weight_decay=0.0008,  # 增加权重衰减防止过拟合
        patience=10,  # 早停机制监视验证损失
        batch=32,  # 减小batch_size适应多尺度训练
        imgsz=768,
        device=0,
        project="runs/detect",
        workers=0,  # Windows下禁用多进程
        cos_lr=True,  # 启用余弦退火学习率
        close_mosaic=20,  # 最后15个epoch关闭Mosaic
        optimizer='AdamW',  # 更换优化器
        mixup=0.5,  # 覆盖配置中的MixUp概率
        copy_paste=0.2,  # 覆盖配置中的copy-paste概率
        erasing=0.6,  # 显式启用随机擦除增强
        scale=0.5  # 多尺度训练的初始缩放因子
    )